`include "defines.v"

module EX_Stage (
    //addr
    input  wire [63: 0] pc_i,
    input  wire [ 4: 0] rd_addr_i,
    input  wire [11: 0] csr_addr_i,
    
    output wire [63: 0] pc_o,
    output wire [ 4: 0] rd_addr_o,
    output wire [11: 0] csr_addr_o,

    //data
    input  wire [`xlen] regbusA,
    input wire          rs1_data_sign,
    input  wire [`xlen] regbusB,
    input  wire [`xlen] imm,
    output wire [`xlen] ALU_out,
    output wire [`xlen] data_to_mem,

    //control sigs
    input  wire [ 2: 0] branch_i,
    input  wire         jump_i,
    input  wire         adder_oprendA_src,
    input  wire         is_system_inst_i,
    output wire         is_system_inst_o,
    output wire [ 2: 0] branch_o,
    output wire         jump_o,

    input  wire [ 1: 0] ALU_oprendA_src,
    input  wire [ 1: 0] ALU_oprendB_src,
    input  wire [ 2: 0] ALU_op,
    input  wire [ 2: 0] ALU_out_ext_type,

    input  wire [ 3: 0] MemWr_i,
    input  wire         MemToReg_i,
    input  wire         RegWr_i,
    input  wire [ 2: 0] load_ext_type_i,
    output wire [ 3: 0] MemWr_o,
    output wire         MemToReg_o,
    output wire         RegWr_o,
    output wire [ 2: 0] load_ext_type_o,

    output wire [ 3: 0] flags
);

//pc addr calc
wire [63: 0] adder_oprendA;
wire [63: 0] adder_oprendB;

assign adder_oprendA = adder_oprendA_src ? regbusA  : pc_i;
assign adder_oprendB = imm;
assign pc_o = adder_oprendA + adder_oprendB;
//

//ALU
reg         [`xlen] ALU_oprendA;
reg         [`xlen] ALU_oprendB;
reg         [64: 0] ALU_out_ext;
wire        [`xlen] ALU_out_cut;

always @(*) begin
    case (ALU_oprendA_src)
        `pc              : begin
            ALU_oprendA = pc_i;
        end 
        `regbusA         : begin
            ALU_oprendA = regbusA;
        end
        `zext_bitsel_rs1 : begin
            ALU_oprendA = {32'b0,regbusA[31:0]};
        end
        `sext_bitsel_rs1 : begin
            ALU_oprendA = { {32{rs1_data_sign}} ,regbusA[31:0]};
        end
        default : ALU_oprendA = regbusA;
    endcase
end

always @(*) begin
    case (ALU_oprendB_src)
        `regbusB         : begin
            ALU_oprendB = regbusB;
        end 
        `immbus          : begin
            ALU_oprendB = imm;
        end
        `const4          : begin
            ALU_oprendB = 64'b0100;
        end
        `zext_bitsel_rs2 : begin
            ALU_oprendB = ALU_out_ext_type[0] ? {59'b0,regbusB[4:0]} : {58'b0,regbusB[5:0]};
        end
        default : ALU_oprendB = regbusB;
    endcase
end

assign data_to_mem = regbusB;

always @(*) begin
    case (ALU_op)
        `add : begin
            ALU_out_ext = { 1'b0 , ALU_oprendA } + { 1'b0 , ALU_oprendB };
        end
        `sub : begin
            ALU_out_ext = { 1'b0 , ALU_oprendA } - { 1'b0 , ALU_oprendB };
        end
        `and : begin
            ALU_out_ext = { 1'b0 , ALU_oprendA } & { 1'b0 , ALU_oprendB };
        end
        `or  : begin
            ALU_out_ext = { 1'b0 , ALU_oprendA } | { 1'b0 , ALU_oprendB };
        end
        `xor : begin
            ALU_out_ext = { 1'b0 , ALU_oprendA } ^ { 1'b0 , ALU_oprendB };
        end
        `sll : begin
            ALU_out_ext = { 1'b0 , ALU_oprendA } << { 1'b0 , ALU_oprendB };
        end
        `srl : begin
            ALU_out_ext = { 1'b0 , ALU_oprendA } >> { 1'b0 , ALU_oprendB };
        end
        `sra : begin
            ALU_out_ext = $signed({ rs1_data_sign , ALU_oprendA }) >>> { 1'b0 , ALU_oprendB };
        end
        default : ALU_out_ext = { 1'b0 , ALU_oprendA } + { 1'b0 , ALU_oprendB };
    endcase
end

assign ALU_out_cut = ALU_out_ext[63:0];
assign ALU_out = ALU_out_ext_type[2] ? (ALU_out_ext_type[1] ? {63'b0,`CF} : {63'b0,`SF^`OF} )      :
                 ALU_out_ext_type[0] ? { {32{ALU_out_cut[31]}} , ALU_out_cut[31:0] } : ALU_out_cut ;

assign `CF = ALU_out_ext[64];
assign `SF = ALU_out_ext[63];
assign `OF = (ALU_oprendA[63] & ~ALU_oprendB[63] & ~ALU_out_ext[63])|(~ALU_oprendA[63] & ALU_oprendB[63] & ALU_out_ext[63]);
assign `ZF = ~(|ALU_out_ext);
//

//control sigs
assign rd_addr_o = rd_addr_i;

assign branch_o = branch_i;
assign jump_o = jump_i;

assign MemWr_o = MemWr_i;
assign MemToReg_o = MemToReg_i;
assign RegWr_o = RegWr_i;
assign load_ext_type_o = load_ext_type_i;

assign is_system_inst_o = is_system_inst_i;
//
    
endmodule